import torch
import math
import json
import matplotlib.pyplot as plt
import numpy as np
import math
from random import random
class EdgeNode:
    """
    表示一条边，包含长度、方向、起点、终点，以及辅助的中心点。
    新增：start_node_id和end_node_id记录对应的node列表中的索引。
    """
    def __init__(self, start, end, start_node_id=-1, end_node_id=-1):
        self.start = start  # 起点 (x, y)
        self.end = end  # 终点 (x, y)
        self.length = self.calculate_length(start, end)  # 边的长度
        self.direction = self.calculate_angle(start, end)  # 边的方向
        self.center = ((start[0] + end[0]) / 2, (start[1] + end[1]) / 2)  # 边的中心点
        self.start_node_id = start_node_id
        self.end_node_id = end_node_id

    @staticmethod
    def calculate_length(start, end):
        """计算边的长度"""
        return math.sqrt((end[0] - start[0]) ** 2 + (end[1] - start[1]) ** 2)

    @staticmethod
    def calculate_angle(start, end):
        """计算边的方向（角度，单位：度）"""
        return math.degrees(math.atan2(end[1] - start[1], end[0] - start[0]))


class Graph:
    """
    图结构，存储节点和边，以及图像的尺寸信息。
    """
    def __init__(self, h, w,nodes=[],edges=[]):
        self.nodes = nodes  # 节点集合，每个节点为 (x, y)
        self.edges = edges  # 边集合，每个边为 EdgeNode 对象
        self.h = h  # 原始图像高度
        self.w = w  # 原始图像宽度

    def _2json(self):
        """将图结构保存为 JSON 文件"""
        data = {
            "nodes": self.nodes,
            "edges": [{
                "start": edge.start,
                "end": edge.end,
                "length": edge.length,
                "direction": edge.direction,
                "start_node_id": edge.start_node_id,
                "end_node_id": edge.end_node_id
            } for edge in self.edges],
            "h": self.h,
            "w": self.w
        }
        return data

    def load(self, data):
        """从 JSON 文件加载图结构"""
        self.nodes = [tuple(node) for node in data["nodes"]]
        self.edges = []
        for e in data["edges"]:
            edge = EdgeNode(tuple(e["start"]), tuple(e["end"]), e["start_node_id"], e["end_node_id"])
            self.edges.append(edge)
        self.h = data["h"]
        self.w = data["w"]

    def build_from_vessel(self, vessels):
        """
        从给定的血管列表构建图结构。
        :param vessels: 由 VesselNode 组成的血管列表。
                        假设每个vessel拥有属性nodes，其中每个node有start和end属性。
        """
        for vessel in vessels:
            for node in vessel.nodes:
                start_idx = self.add_or_get_node(node.start)
                end_idx = self.add_or_get_node(node.end)
                start_coord = self.nodes[start_idx]
                end_coord = self.nodes[end_idx]
                edge = EdgeNode(start_coord, end_coord, start_node_id=start_idx, end_node_id=end_idx)
                self.edges.append(edge)

    def add_or_get_node(self, point):
        """
        检查点是否已存在于节点集合中，若存在则返回已有点的索引；否则加入集合并返回新索引。
        :param point: 待添加的点 (x, y)
        :return: 在节点集合中的点索引
        """
        for i, node in enumerate(self.nodes):
            if self.calculate_distance(node, point) < 2:
                return i
        self.nodes.append(point)
        return len(self.nodes) - 1

    @staticmethod
    def calculate_distance(p1, p2):
        """计算两点之间的欧几里得距离"""
        return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)

    def visualization(self, save_path=None):
        """
        可视化图结构：
        节点用黄色实心圆表示，边用蓝色线段连接。
        """
        plt.figure(figsize=(self.w / 100, self.h / 100))
        plt.axis([0, self.w, 0, self.h])
        
        # 绘制节点
        for node in self.nodes:
            if node != (-1, -1):  # 只绘制未删除的节点
                plt.plot(node[0], node[1], 'yo', markersize=3)

        # 绘制边
        for edge in self.edges:
            if edge.start != (-1, -1) and edge.end != (-1, -1):
                plt.plot([edge.start[0], edge.end[0]], [edge.start[1], edge.end[1]], 'b-')

        plt.gca().invert_yaxis()  # 图像坐标以左上角为原点
        if save_path:
            plt.savefig(save_path)
        else:    
            plt.show()
    
    

    def to_tensors(self, max_length, device='cpu'):
        """
        将graph转换为适合GNN输入的张量：
        返回： 
         node_positions: [N, 2]
         edge_features: [E, input_dim] (对length和direction进行特征工程)
         edge_index: [2, E]
        特征工程:
          - length 归一化：length / max_length
          - direction 转换成 (cos, sin)
        """
        node_positions = torch.tensor(self.nodes, dtype=torch.float, device=device)
        
        edge_features = []
        edge_index = []
        for e in self.edges:
            normalized_length = e.length / max_length
            dir_rad = math.radians(e.direction)
            dir_x = math.cos(dir_rad)
            dir_y = math.sin(dir_rad)
            # 边特征为 [normalized_length, dir_x, dir_y]
            edge_features.append([normalized_length, dir_x, dir_y])
            edge_index.append([e.start_node_id, e.end_node_id])

        edge_features = torch.tensor(edge_features, dtype=torch.float, device=device)
        edge_index = torch.tensor(edge_index, dtype=torch.long, device=device).t()  # [2, E]
        return node_positions, edge_features, edge_index



class GraphTransformer():
    def __init__(self,image_size=(1600,1200),
                 camrea_weight=1570,
                 camera_angle=math.pi/2,
                 sample_dense=10,
                 threshold=0.5):
        self.weight,self.height=image_size
        self.camera_angle=camera_angle
        self.radius=camrea_weight/(2*math.sin(camera_angle/2))
        self.sample_dense=sample_dense
        self.threshold=threshold

        # build_z_matric
        self.z_matric=np.zeros((self.height,self.weight))
        for i in range(self.height):
            for j in range(self.weight):
                x,y=self._get_xy(j,i)
                self.z_matric[i][j]=self._cal_z(x,y)

    def _cal_z(self,x,y):
        tmp=self.radius**2-x**2-y**2
        if tmp>=0:
            return round(math.sqrt(tmp))
        return 0
    def _get_z(self,x,y):
        coor_x,coor_y=self._xy2coor(x,y)
        return self.z_matric[coor_y][coor_x]
    def _get_xy(self,coordinate_x,coordinate_y):
        return coordinate_x-(self.weight/2),-coordinate_y+(self.height/2)
    
    def _xy2coor(self,x,y):
        return x+(self.weight/2),(self.height/2)-y
    
    def calculate_angle(self,ridge_x,ridge_y,optic_x,optic_y):
        z_ridge = self._get_z(ridge_x, ridge_y)
        z_optic = self._get_z(optic_x, optic_y)
        l = math.sqrt((ridge_x - optic_x)**2 + (ridge_y - optic_y)**2 + (z_ridge - z_optic)**2)
        # in circle l ischord length, cal the angle with l and radius
        try:
            angle=2*math.asin(l/(2*self.radius))
        except:
            print(l)
            raise
        return math.degrees(angle)
    def _get_delta(self, src_x, src_y, d_theta, axis='x', low=1, high=100):
        direction = 1 if d_theta > 0 else -1
        while True:
            mid = (low + high) / 2

            # 根据axis决定对哪个坐标进行调整
            if axis == 'x':
                tar_x = src_x + mid * direction
                tar_y = src_y
            else:
                tar_x = src_x
                tar_y = src_y + mid * direction 

            angle = self.calculate_angle(src_x, src_y, tar_x, tar_y)

            if abs(angle - d_theta) < self.threshold:
                return mid
            if angle < d_theta:
                low = mid
            else:
                high = mid
            if high - low < 2:
                return mid  

    def get_dy(self, src_x, src_y, d_theta, low=1, high=100):
        return self._get_delta(src_x, src_y, d_theta, axis='y', low=low, high=high) 

    def _get_dx(self, src_x, src_y, d_theta, low=1, high=100):
        return self._get_delta(src_x, src_y, d_theta, axis='x', low=low, high=high) 
    def noise(e=3):
        sample = np.random.normal(loc=0, scale=1)  # 生成标准正态分布采样
        sample=min(max(sample,-e),e)
        return sample
    def modify_points_list(self,points):
        new_points=[]
        d_x_theta,d_y_theta=self.get_random_d_tha(),self.get_random_d_tha()
        missing_number=0
        for i in range(len(points)-1):
            src_x,src_y=points[i]
            dx=self._get_dx(src_x,src_y,d_x_theta+self.noise())
            dy=self.get_dy(src_x,src_y,d_y_theta+self.noise)
            tar_x,tar_y=src_x+dx,src_y+dy
            if tar_x<0 or tar_x>=self.weight or tar_y<0 or tar_y>=self.height:
                tar_x,tar_y=-1,-1
                missing_number+=1
            new_points.append((tar_x,tar_y))
            
        return new_points,missing_number
    def get_random_d_tha(self):
        # 返回 -45，45之间的随机数
        return int(random()*90-45)

def transform_graph(graph, transformer):
    '''
    对图进行转换，返回新的图对象。
        使用new_points更新所有节点的坐标。
        new_points和self.nodes长度必须一致。
        若new_points中某坐标为(-1,-1)，则表示该节点消失，与该节点相关的边将不保留。
        
        更新完成后，self.nodes = new_points
        self.edges根据保留的节点重建。
    '''
    new_points,missing_number=transformer.modify_points_list(graph.nodes)
        
    if len(graph.nodes) != len(new_points):
        raise ValueError('The length of new points is not equal to the length of nodes')
    # 更新节点坐标
    node_map = []
    new_points_clr = []
    new_idx=0
    for i, node in enumerate(graph.nodes):
        if node != (-1, -1):
            node_map.append(new_idx)
            new_points_clr.append(new_points[i])
            new_idx+=1
        else:
            node_map.append(-1)
            
    
    new_edges = []
    for edge in graph.edges:
        start_id = edge.start_node_id
        end_id = edge.end_node_id
        # 如果任一节点消失，则该边不再保留
        if new_points[start_id] == (-1, -1) or new_points[end_id] == (-1, -1):
            # 用更新前的id查看是否消失
            continue
        
        # 否则使用更新后的坐标重建边
        new_start = new_points[start_id]
        new_end = new_points[end_id]
        new_start_id = node_map[start_id]
        new_end_id = node_map[end_id]
        
        new_edge = EdgeNode(new_start, new_end, start_node_id=new_start_id, end_node_id=new_end_id)
        new_edges.append(new_edge)
    
    
    node_map=graph._build_new_graph(new_points)
    new_graph = Graph(graph.h, graph.w, new_points_clr, new_edges)
    return new_graph,missing_number,node_map

if __name__=='__main__':
    pass